//===- iteration_predictor.cpp - Iteration Predictor Pass" -==================//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements a Iteration Predictor Pass.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "iteration-predictor"
#include "llvm/Pass.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/LoopInfo.h"
#include <llvm/IR/Constants.h>
#include <llvm/IR/DerivedTypes.h>
#include "llvm/IR/Function.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/IRBuilder.h"
#include <llvm/IR/LLVMContext.h>
#include "llvm/IR/Module.h"
#include "llvm/Support/TargetFolder.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetLibraryInfo.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"

#include <vector>

using namespace llvm;

STATISTIC(FunctionsAnalyzed, "Number of functions in the program");
STATISTIC(LoopsAnalyzed,     "Number of loops analyzed");


namespace {
  // IterationPredictor 
  struct IterationPredictor : public FunctionPass {
    static char ID; // Pass identification, replacement for typeid

    LLVMContext *context;
    GlobalVariable* GVstderr;
    Value* SubFPrintf;
    Value* IncFPrintf;
    Value* EndFPrintf;
    Module* TheModule;
    IterationPredictor() : FunctionPass(ID), context(NULL), 
                           GVstderr(NULL), SubFPrintf(NULL), 
                           IncFPrintf(NULL), EndFPrintf(NULL), 
                           TheModule(NULL) {}

    void init(Module* M) {
      this->TheModule = M;
      this->context = &TheModule->getContext();
    }

    void createFPrintf(Function& F) {
      Module* M = F.getParent();
      Type* IO_FILE_PTR_ty;
      GVstderr = TheModule->getGlobalVariable("stderr");
      if (GVstderr != NULL) {
        IO_FILE_PTR_ty = GVstderr->getType();
        IO_FILE_PTR_ty = IO_FILE_PTR_ty->getContainedType(0);
      } else {
        StructType* IO_FILE_ty = StructType::create(*context, 
                                                    "struct._IO_FILE");
        IO_FILE_PTR_ty = PointerType::getUnqual(IO_FILE_ty);
        StructType* IO_marker_ty = StructType::create(*context, 
                                                      "struct._IO_marker");
        PointerType* IO_marker_ptr_ty = PointerType::getUnqual(IO_marker_ty);

        std::vector<Type*> Elements;
        Elements.push_back(Type::getInt32Ty(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(IO_marker_ptr_ty);
        Elements.push_back(IO_FILE_PTR_ty);
        Elements.push_back(Type::getInt32Ty(*context));
        Elements.push_back(Type::getInt32Ty(*context));
        Elements.push_back(Type::getInt32Ty(*context));
        Elements.push_back(Type::getInt16Ty(*context));
        Elements.push_back(Type::getInt8Ty(*context));
        Elements.push_back(ArrayType::get(Type::getInt8Ty(*context), 1));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt64Ty(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt8PtrTy(*context));
        Elements.push_back(Type::getInt32Ty(*context));
        Elements.push_back(Type::getInt32Ty(*context));
        Elements.push_back(ArrayType::get(Type::getInt8Ty(*context), 40));
        IO_FILE_ty->setBody(Elements, false);

        std::vector<Type*> Elements2;
        Elements2.push_back(IO_marker_ptr_ty);
        Elements2.push_back(IO_FILE_PTR_ty);
        Elements2.push_back(Type::getInt32Ty(*context));
        IO_marker_ty->setBody(Elements2, false);
        GVstderr = new GlobalVariable(*M, IO_FILE_PTR_ty, false,
                                      GlobalValue::ExternalLinkage, NULL, 
                                      "stderr");
      }

      std::vector<Type*> Params;      
      Params.push_back(IO_FILE_PTR_ty);
      Params.push_back(PointerType::getUnqual(Type::getInt8Ty(*context)));
      Params.push_back(Type::getInt32Ty(*context));

      Type* Ty = Type::getVoidTy(*context);
      FunctionType *T = FunctionType::get(Ty, Params, true);
      SubFPrintf = TheModule->getOrInsertFunction("fprintf", T);
      IncFPrintf = TheModule->getOrInsertFunction("fprintf", T);
      EndFPrintf = TheModule->getOrInsertFunction("fprintf", T);

    }


    Value* insertAdd(BasicBlock *BB, AllocaInst *A) {
      IRBuilder<> Builder(BB->getFirstNonPHI());
      Value* L = Builder.CreateAlignedLoad(A, 4);
      ConstantInt* One = ConstantInt::get(Type::getInt32Ty(*context), 1);
      Value* I = Builder.CreateAdd(L, One);
      Builder.CreateAlignedStore(I, A, 4);
      return I;
    }


    AllocaInst* insertAlloca(BasicBlock *BB, Constant* constZero) {
      Type* Ty = IntegerType::get(BB->getParent()->getContext(), 32);
      AllocaInst* A = new AllocaInst(Ty, "", BB->getFirstNonPHI());
      StoreInst* store = new StoreInst(constZero, A, false, 4, 
                                       BB->getTerminator());
      return A;
    }


    void InsertPrintInsts(BasicBlock* BB, Value* V, Value* subStr, 
                          AllocaInst* A, Instruction* InsPt) {

      IRBuilder<> Builder(BB->getFirstInsertionPt());

      // Store 1, V
      ConstantInt* One = ConstantInt::get(Type::getInt32Ty(*context), 1000);
      Builder.CreateAlignedStore(One, V, 4);

      Value* Sub = Builder.CreateSub(InsPt->getOperand(0), InsPt->getOperand(1));
      Builder.CreateAlignedStore(Sub, A, 4);

      // Print sub.
      Value* L = Builder.CreateAlignedLoad(GVstderr, 4);
      std::vector<Value*> args;
      args.push_back(L);
      args.push_back(subStr);
      args.push_back(Sub);
      Builder.CreateCall(SubFPrintf, args, "");
    }


    BasicBlock* createBlock(BasicBlock *Old, Value *V, Value *Z,
                             AllocaInst* A, Value* subStr) {
      Instruction* LInsPt = NULL;
      // Find the last branch in the block.
      for (BasicBlock::iterator I = Old->begin(), E = Old->end(); I != E; ++I) {
        if (CmpInst *Cmp = dyn_cast<CmpInst>(I)) {
          LInsPt = Cmp;
        }
      }

      Type* Ty = Type::getInt32Ty(*context);
      
      if (LInsPt == NULL)
        return NULL;

      if (LInsPt->getOperand(0)->getType() != Ty || LInsPt->getOperand(1)->getType() != Ty)
        return NULL;

      // Load the variable that controls the printing and compare it to zero.
      IRBuilder<> Builder(LInsPt);
      Value* L = Builder.CreateAlignedLoad(V, 4);
      Builder.CreateICmpEQ(L, Z);

      // Find the first branch in the block.
      Instruction* FInsPt = NULL;
      for (BasicBlock::iterator I = Old->begin(), E = Old->end(); I != E; ++I) {
        if (CmpInst *Cmp = dyn_cast<CmpInst>(I)) {
          FInsPt = Cmp;
          break;
        }
      }

      TerminatorInst* Tr = SplitBlockAndInsertIfThen(FInsPt, false);
      BasicBlock* newBB = Tr->getParent();
      InsertPrintInsts(newBB, V, subStr, A, LInsPt);
      return newBB;
    }

    virtual bool runOnFunction(Function &F) {
      init(F.getParent());
      ConstantInt* constZero = ConstantInt::get(Type::getInt32Ty(*context), 0);
      FunctionsAnalyzed++;
      IRBuilder<> Builder(&F.getEntryBlock());
      Value* subStr = Builder.CreateGlobalStringPtr("S: %d ",    "sub_str");
      Value* incStr = Builder.CreateGlobalStringPtr("I: %d\n",   "inc_str");
      Value* endStr = Builder.CreateGlobalStringPtr("====\n",    "end_str");

      createFPrintf(F);
      LoopInfo &LI = getAnalysis<LoopInfo>(); 
      for (LoopInfo::iterator I = LI.begin(), E = LI.end(); I != E; ++I) {
        if ((*I)->getLoopDepth() != 1)
          continue;

        AllocaInst *Sub = insertAlloca(&F.getEntryBlock(), constZero);
        AllocaInst *Inc = insertAlloca(&F.getEntryBlock(), constZero);
        AllocaInst *V   = insertAlloca(&F.getEntryBlock(), constZero);

        /// A latch block is a block that contains a branch back to the header.
        BasicBlock *Latch  = (*I)->getLoopLatch();
        BasicBlock *Header = (*I)->getHeader();
        BasicBlock* FBlock = createBlock(Latch, V, constZero, Sub, subStr);
        if (FBlock == NULL)
          FBlock = createBlock(Header, V, constZero, Sub, subStr);

        if (FBlock == NULL)
          continue;        

        Value* A = insertAdd(Header, Inc);
        assert(A != NULL && "Couldn't insert 'add'!");

        const BranchInst *FBr = dyn_cast<BranchInst>(FBlock->getTerminator());
        BasicBlock *OldLatch = FBr->getSuccessor(0);
        const BranchInst *Br = dyn_cast<BranchInst>(OldLatch->getTerminator());
        assert(Br && Br->isConditional());
        
        BasicBlock *FalseBlock = Br->getSuccessor(1);
        assert(FalseBlock != NULL && "No False Block!");

        IRBuilder<> Builder(FalseBlock->getFirstNonPHI());
        Value* L = Builder.CreateAlignedLoad(GVstderr, 4);
        std::vector<Value*> args;
        args.push_back(L);
        args.push_back(incStr);
        args.push_back(A);
        Builder.CreateCall(IncFPrintf, args, "");

        args.clear();
        args.push_back(L);
        args.push_back(endStr);
        args.push_back(A);
        Builder.CreateCall(EndFPrintf, args, "");

        LoopsAnalyzed++;
        return true;
      }
      
      return false;
    }


    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.addRequired<LoopInfo>();
    }  
  };
}

char IterationPredictor::ID = 0;
static RegisterPass<IterationPredictor>
Y("iteration-predictor", "Iteration Predictor Pass");
